/**
 * 
 */
package edu.arizona.bio5.onto.decomposition.persistent.xml;

import java.io.IOException;
import java.io.Reader;
import java.io.StringReader;
import java.net.URI;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import org.semanticweb.owlapi.apibinding.OWLManager;
import org.semanticweb.owlapi.io.StringDocumentSource;
import org.semanticweb.owlapi.model.OWLAxiom;
import org.semanticweb.owlapi.model.OWLOntology;
import org.semanticweb.owlapi.model.OWLOntologyCreationException;
import org.semanticweb.owlapi.model.OWLOntologyManager;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.NodeList;

import uk.ac.manchester.cs.graph.DirectedGraph;
import uk.ac.manchester.cs.graph.Graph;
import uk.ac.manchester.cs.graph.io.GraphMLReaderImpl;
import uk.ac.manchester.cs.graph.io.GraphReader;
import uk.ac.manchester.cs.owlapi.modularity.ModuleType;
import edu.arizona.bio5.onto.decomposition.Atom;
import edu.arizona.bio5.onto.decomposition.persistent.ADInfo;
import edu.arizona.bio5.onto.decomposition.persistent.ADLoader;
import edu.arizona.bio5.onto.decomposition.persistent.ADPersistenceException;
import edu.arizona.bio5.onto.decomposition.persistent.AtomLoader;
import edu.arizona.bio5.onto.decomposition.persistent.PersistentAtomReader;
import edu.arizona.bio5.onto.decomposition.persistent.PersistentAtomicDecomposition;
import edu.arizona.bio5.onto.decomposition.persistent.PersistentAtomicDecompositionImpl;

/**
 * @author Pavel Klinov
 *
 */
public abstract class XMLADLoader implements ADLoader {

	@Override
	public PersistentAtomicDecomposition load() throws ADPersistenceException {
		prepare();
		
		try {
			ADInfo adInfo = loadADInfo();
			Map<Object, Atom> idAtomMap = new HashMap<Object, Atom>();
			DirectedGraph<Atom, Object> adGraph = loadGraph(idAtomMap, createXMLAtomLoader());
			
			loadLabels(idAtomMap, new XMLSeedSigLabelLoader(OWLManager.getOWLDataFactory()));
			
			Map<URI, Set<Atom>> tbIndex = loadTermBasedIndex(idAtomMap);
			Set<OWLAxiom> globalAxioms = loadGlobalAxioms();

			PersistentAtomicDecompositionImpl pda = new PersistentAtomicDecompositionImpl(URI.create(adInfo.getOntologyURI()), adGraph, tbIndex, adInfo.getModType());
			
			pda.setGlobalAxioms(globalAxioms);
			pda.setTermBasedIndex(tbIndex);
			
			return pda;
		} catch (Throwable e) {
			throw e instanceof ADPersistenceException ? (ADPersistenceException)e : new ADPersistenceException(e);
		} finally {
			finish();
		}
	}

	private Set<OWLAxiom> loadGlobalAxioms() throws ADPersistenceException {
		String gaDoc = loadDocumentAsString(XMLADConstants.GLOBAL_AXIOMS_DOC);
		OWLOntologyManager mgr = OWLManager.createOWLOntologyManager();

		try {
			OWLOntology gaOntology = mgr.loadOntologyFromOntologyDocument(new StringDocumentSource(gaDoc));
			
			return gaOntology.getAxioms();
		} catch (OWLOntologyCreationException e) {
			throw new ADPersistenceException(e);
		}
	}

	private Map<URI, Set<Atom>> loadTermBasedIndex(Map<Object, Atom> idAtomMap) throws ADPersistenceException {
		Document indexDoc = loadDocument(XMLADConstants.TERM_INDEX_DOC);
		NodeList termElems = indexDoc.getElementsByTagName(XMLADConstants.TERM_ELEM);
		Map<URI, Set<Atom>> tbIndex = new HashMap<URI, Set<Atom>>();
		
		for (int i = 0; i < termElems.getLength(); i++) {
			Element termElem = (Element)termElems.item(i);
			URI term = URI.create(termElem.getAttribute(XMLADConstants.TERM_URI_ATTRIBUTE));
			NodeList atomElems = termElem.getElementsByTagName(XMLADConstants.ATOM_ELEM);
			
			for (int j = 0; j < atomElems.getLength(); j++) {
				Element atomElem = (Element)atomElems.item(j);
				Atom atom = idAtomMap.get(atomElem.getAttribute(XMLADConstants.ATOM_ID_ATTR));
				Set<Atom> termAtoms = tbIndex.get(term);
				
				termAtoms = termAtoms == null ? new HashSet<Atom>() : termAtoms;
				termAtoms.add(atom);
				tbIndex.put(term, termAtoms);
			}
		}
		
		return tbIndex;
	}

	private void loadLabels(Map<Object, Atom> idAtomMap, XMLSeedSigLabelLoader labelLoader) throws ADPersistenceException {
		Document labelsDoc = loadDocument(XMLADConstants.LABELS_DOC);
		NodeList atomElems = labelsDoc.getElementsByTagName(XMLADConstants.ATOM_ELEM);
		
		for (int i = 0; i < atomElems.getLength(); i++) {
			Element atomElem = (Element)atomElems.item(i);
			String atomID = atomElem.getAttribute(XMLADConstants.ATOM_ID_ATTR);
			Atom atom = idAtomMap.get(atomID);
			
			if (atom != null) {
				//Load the label for the atom
				NodeList labelNodes = atomElem.getElementsByTagName(XMLADConstants.LABEL_ELEM);
				
				for (int j = 0; j < labelNodes.getLength(); j++) {
					Element labelElement = (Element)labelNodes.item(j);
					
					labelLoader.setInput(labelElement);
					atom.setLabel(labelLoader.load());
				}				
			}
		}
	}

	protected ADInfo loadADInfo() throws ADPersistenceException {
		Document adInfoDoc = loadDocument(XMLADConstants.AD_INFO_DOC);
		Element root = adInfoDoc.getDocumentElement();
		ModuleType modType = ModuleType.valueOf(root.getAttribute(XMLADConstants.AD_MODULE_TYPE_ATTR));
		Element uriElem = (Element)root.getElementsByTagName(XMLADConstants.AD_ONTO_URI_ELEMENT).item(0);
		
		return new ADInfo(uriElem.getTextContent(), modType);
	}

	protected DirectedGraph<Atom, Object> loadGraph(Map<Object, Atom> idAtomMap, AtomLoader atomLoader)  throws ADPersistenceException {
		//TODO It might be better to use XML:DB API directly, e.g., getContentAsSAX
		Reader docReader = new StringReader(loadDocumentAsString(XMLADConstants.GRAPH_DOC));
		GraphReader<Atom, Object> reader = new GraphMLReaderImpl<Atom, Object>();
		
		reader.setNodeReader(new PersistentAtomReader(atomLoader));
		
		Graph<Atom, Object> adGraph;
		try {
			adGraph = reader.read(docReader);
		} catch (IOException e) {
			throw new ADPersistenceException(e);
		}
		
		for (Atom atom : adGraph.getVertices()) idAtomMap.put(atom.getID(), atom);
		
		return (DirectedGraph<Atom, Object>)adGraph;
	}
	
	protected abstract void prepare() throws ADPersistenceException;
	protected abstract void finish();
	protected abstract Document loadDocument(String docName) throws ADPersistenceException;
	protected abstract String loadDocumentAsString(String docName) throws ADPersistenceException;
	protected abstract AtomLoader createXMLAtomLoader() throws ADPersistenceException;

}